import numpy as np
from anytree import AnyNode,search,LevelOrderIter,RenderTree,Node,Walker
import copy
import math
import sys
import os
sys.path.append('../')
from utility.keras_modal import DenseModel
from board_fast import Board,Game,Player,Point,Move,GameResult
from utility.hdf5 import HDF5

class agz_tree:
    def __init__(self,player):
        self.root=AnyNode(n=1,player=player)
    def getNodeLeaves(self,node_start): #获取子节点
        return [node for node in LevelOrderIter(node_start, filter_=lambda n: n is not node_start,maxlevel=2)]
    def addNode(self,father,kid,player):
        Np=father.n
        c=2
        P=kid['item']
        if player!=self.root.player:
            P=-P
        n=0
        Q=0
        T=Q+c*P*(math.sqrt(Np)/(1+n))
        kid=AnyNode(parent=father,n=n,P=P,Q=Q,T=T,move=kid['move'],player=player)
    def findNextMove(self,node):
        nodes=self.getNodeLeaves(node)
        #print(RenderTree(node))
        if len(nodes)>1:   #仅有一个节点的话，说明没有可以落子的地方了
            node=max(nodes,key=lambda k:abs(k.T))
            node.n+=1
            c=2
            Np=node.parent.n
            node.T=node.Q/node.n+c*node.P*(math.sqrt(Np)/(1+node.n))
            return node
        else:
            return None
    def findMove(self):
        nodes=self.getNodeLeaves(self.root)
        if len(nodes)!=0:
            node=max(nodes,key=lambda k:k.n)
            return node.move
        else:
            return None


    def update_nodes(self,node,root,winner):
        w=Walker()
        nodes=w.walk(node,root)
        rollout=0
        c=2
        for i in nodes[0]:
            if i is not root:
                if winner is not None:
                    if i.player==winner:
                        rollout=1
                    elif i.player.other()==winner:
                        rollout=-1
                    else:
                        None
                else:
                    rollout=0
                i.Q+=rollout
                V=i.Q/i.n
                Np=i.parent.n
                i.T=V+c*i.P*(math.sqrt(Np)/(1+i.n))

class ai_bot:
    def __init__(self,boardSize=9,rand=False):
        self.model=DenseModel(boardSize=boardSize,model='AGZ')
        self.boardSize=boardSize
        self.moves=frozenset()

    def load_weights(self,fileName):
        self.model.model_load_weights(fileName)
    def unload_weights(self,fileName):
        self.model.model_save_weights(fileName)

    def compile(self):
        self.model.compile_ac()
    def train(self,x,y):
        y_train_p=y[:,:-1]
        y_train_v=y[:,-1:]
        self.model.fit_ac(x,y_train_p,y_train_v,batch_size=32,epochs=1)

    def isLegal(self,point,board,player):
        neighbours=board.get_neighbors(point)
        is_eye=True
        is_self_captured=True
        the_other_strings=[]
        own_strings=[]
        liberties=[]
        #判断是不是眼位
        for i in neighbours:
            if board.grid.get(i)==None:
                is_eye=False
                break
            elif board.grid.get(i).player != player:
                is_eye=False
                break
            elif len(board.grid.get(i).liberties)<=1:
                is_eye=False
                break
            else:
                pass
        #判断是不是自杀
        for i in neighbours:
            go_string_t=board.grid.get(i)
            if go_string_t is not None: #有人占据了
                if go_string_t.player == player:    #自己的
                    if go_string_t not in own_strings:
                        own_strings.append(go_string_t)
                else:
                    if go_string_t not in the_other_strings:
                        the_other_strings.append(go_string_t)
            else:   #空子
                liberties.append(i)
        if len(liberties) == 0:
            if any(len(str.liberties) == 1 for str in the_other_strings):   #如果落子前对方有任一串棋只有一口气了，就是提子
                is_self_captured=False
            elif all(len(str.liberties) == 1 for str in own_strings):
                None    #自杀
            else:
                is_self_captured=False
        else:
            is_self_captured=False
        if is_eye or is_self_captured:
            return False
        return True


    def chooseMove(self,player,board,reChoose=False,mct_rounds=1600,depth=180,isTrain=False):
        if not reChoose:
            self.moves=frozenset()
        root_tree=agz_tree(player)
        boardsize=board.size
        value_first=None
        for _ in range(mct_rounds):
            board_now=copy.deepcopy(board)
            player_now=player
            node_now=root_tree.root
            node_now.n+=1
            pass_count=0    #由于没有引入game类，需要记录连续两次pass
            for j in range(depth):  #限制陷入全局同形
                #print(j)
                isOver=False
                npboard=board_now.print_board(isprint=False)
                npboard=npboard.reshape(*npboard.shape,1)
                if player_now==Player.black:
                    npplayer=np.ones((boardsize,boardsize))
                else:
                    npplayer=-1*np.ones((boardsize,boardsize))
                npplayer=npplayer.reshape(*npplayer.shape,1)
                nparray=np.concatenate((npboard,npplayer),axis=-1)
                nparray=nparray.reshape(1,*nparray.shape)
                mask_now=(npboard==0).astype(int)
                while True: #直到找到合理的落子，或者没有合理落子游戏结束
                    predict,value=self.model.model_predict_ac(nparray)
                    predict=predict[0]
                    value=value[0]
                    if j == 0:
                        value_first=value   #记录网络的v输出，这是网络学习时要提供的
                        for i in self.moves:    #屏蔽掉已经下过的子
                            mask_now[i]=False
                    pred_valid=mask_now.flatten()*predict[:-1]
                    nodes=root_tree.getNodeLeaves(node_now)
                    if len(nodes)==0:   #没有子节点，就要新增
                        root_tree.addNode(node_now,{'move':(boardsize,boardsize),'item':pred_valid[-1]},player_now) #加上弃手
                        for idx,item in enumerate(pred_valid):
                            if item !=0:    #排除以被占的点位
                                move=np.unravel_index(idx, (boardsize,boardsize))   #move是普通tuple
                                point=Point(move[0]+1,move[1]+1)
                                if self.isLegal(point,board_now,player_now):
                                    root_tree.addNode(node_now,{'move':move,'item':item},player_now)
                    node_=root_tree.findNextMove(node_now)  #因为有弃手，所以node_不会为空
                    if node_ is None:
                        assert(True == False)
                        isOver=True
                        break
                    elif node_.move == (boardsize,boardsize):  #弃手
                        #print("PASS")
                        pass_count+=1

                        if pass_count==2:   #连着两次弃手了
                            node_now=node_
                            move=node_now.move
                            point=Point(move[0]+1,move[1]+1)
                            for i in range(2):  #仿真时连续两手不能结束游戏，回到两手Pass前的状态
                                node_now=node_now.parent
                            break

                        node_now=node_
                        move=node_now.move
                        point=Point(move[0]+1,move[1]+1)
                        break
                    else: #选最大值落子,并走一步
                        pass_count=0
                        node_now=node_
                        move=node_now.move
                        if move is not None:
                            point=Point(move[0]+1,move[1]+1)
                            break
                        else:   #这步逻辑应该是走不到了
                            isOver=True
                            break
                if isOver==True:    #进入结算
                    result=board_now.getGameResult()
                    if result>0:
                        winner=Player.black
                    elif result<0:
                        winner=Player.white
                    else:
                        winner=None
                    root_tree.update_nodes(node_now,root_tree.root,winner)
                    break
                else:
                    if point!=Point(boardsize+1,boardsize+1):
                        board_now.play_stone(player_now,point)
                        player_now=player_now.other()
                    else:
                        player_now=player_now.other()
            if isOver!=True:    #走到了depth参数限制都没走完，就认为是平局
                winner=None
                root_tree.update_nodes(node_now,root_tree.root,winner)
        move=root_tree.findMove()
        if move is not None and move !=(boardsize,boardsize):
            self.moves=self.moves|{move}
            point=Point(move[0]+1,move[1]+1)
            move_=Move(point=point)
        else:
            move_=Move(is_pass=True)
        nodes=root_tree.getNodeLeaves(root_tree.root)
        npdata=np.zeros((boardsize*boardsize+1,))
        npdata_=npdata[:-1].reshape(boardsize,boardsize)
        for i in nodes: #输出待
            if i.move!=(boardsize,boardsize):
                npdata_[i.move]=i.n
            else:
                npdata[-1]=i.n
        return move_,value_first,npdata


class f5obj:    #这个结构和在ac算法里实现的保存几乎是一样的
    def __init__(self,fileName):
        self.fileName=fileName
    def save_game(self,winner,game_recorder,boardSize=9):
        h5file=HDF5(self.fileName,mode='a')
        #grpName=hashlib.new('md5',str(moves_his).encode(encoding='utf-8')).hexdigest()
        grpName='AlphaZero' #AlphaZero的算法如果不学一次更新一次，那么出来的结果是一样的，所以每次只要保存一局
        h5file.add_group(grpName)
        h5file.set_grp_attrs(grpName,'winner',winner)
        h5file.set_grp_attrs(grpName,'size',boardSize)
        for i,dset in enumerate(game_recorder):
            #直接处理好样本和标签，以免后续去值时再做处理
            player=dset[0]
            value_=dset[1]
            if player==winner:
                value=1-value_
            else:
                value=-1-value_
            game_board=dset[2]
            if player=='b':
                player_board=np.ones(game_board.shape)
            else:
                player_board=-1*np.ones(game_board.shape)
            #构建样本输入
            game_board=game_board.reshape(*game_board.shape,1)
            player_board=player_board.reshape(*player_board.shape,1)
            npboard=np.concatenate((game_board,player_board),axis=-1)
            target=dset[3]/dset[3].sum()
            h5file.add_dataset(grpName,str(i),npboard)
            h5file.set_dset_attrs(grpName,str(i),'player',player)
            h5file.set_dset_attrs(grpName,str(i),'target',target)
        h5file.closeH5()

    def load_data(self):
        h5file=HDF5(self.fileName,mode='r')
        x_set=None
        y_set=None
        for i in h5file.grps:
            size=int(h5file.get_grp_attrs(i,'size'))
            winner=h5file.get_grp_attrs(i,'winner')
            for k in h5file.f[i]:
                npboard=np.array(h5file.f[i][k])
                target=h5file.get_dset_attrs(i,k,'target')
                player=h5file.get_dset_attrs(i,k,'player')
                #value=np.array(self.get_dset_attrs(i,k,'value')
                board_target=target[:-1].reshape(size,size)
                x0=npboard
                y0=board_target
                x1=np.rot90(x0, 1)
                y1=np.rot90(y0, 1)
                x2=np.rot90(x0, 2)
                y2=np.rot90(y0, 2)
                x3=np.rot90(x0, 3)
                y3=np.rot90(y0, 3)
                x4=np.flip(x0)
                y4=np.flip(y0)
                x5=np.rot90(x4, 1)
                y5=np.rot90(y4, 1)
                x6=np.rot90(x4, 2)
                y6=np.rot90(y4, 2)
                x7=np.rot90(x4, 3)
                x0=x0.reshape(1,*x0.shape)
                x1=x1.reshape(1,*x1.shape)
                x2=x2.reshape(1,*x2.shape)
                x3=x3.reshape(1,*x3.shape)
                x4=x4.reshape(1,*x4.shape)
                x5=x5.reshape(1,*x5.shape)
                x6=x6.reshape(1,*x6.shape)
                x7=x7.reshape(1,*x7.shape)
                y7=np.rot90(y4, 3)
                y0=y0.flatten()
                y1=y1.flatten()
                y2=y2.flatten()
                y3=y3.flatten()
                y4=y4.flatten()
                y5=y5.flatten()
                y6=y6.flatten()
                y7=y7.flatten()
                if player==winner:
                    value=1
                else:
                    value=-1
                y0=np.hstack((y0,target[-1],value))
                y1=np.hstack((y1,target[-1],value))
                y2=np.hstack((y2,target[-1],value))
                y3=np.hstack((y3,target[-1],value))
                y4=np.hstack((y4,target[-1],value))
                y5=np.hstack((y5,target[-1],value))
                y6=np.hstack((y6,target[-1],value))
                y7=np.hstack((y7,target[-1],value))

                x=np.concatenate((x0,x1,x2,x3,x4,x5,x6,x7),0)

                y=np.vstack((y0,y1,y2,y3,y4,y5,y6,y7))
                if x_set is None or y_set is None:
                    x_set=x
                    y_set=y
                else:
                    x_set=np.vstack((x_set,x))
                    y_set=np.vstack((y_set,y))
        return x_set,y_set




if __name__ == "__main__":
    #board = Board(size=4)
    #game=Game(board)
    #board.print_board(isprint=True)
    weights="./weights/weights.h5"
    hdf5="./hdf5/train.h5"
    bot1=ai_bot(boardSize=4)
    bot2=ai_bot(boardSize=4)
    #if pathlib.Path(ai_weights_s).exists() and pathlib.Path(ai_weights_c).exists() and pathlib.Path(ai_weights_v).exists():
    #    bot1.load_weights(ai_weights_s,ai_weights_c,ai_weights_v)
    #    bot2.load_weights(ai_weights_s,ai_weights_c,ai_weights_v)
    #player=Player.black
    #move,_=bot1.chooseMove(player,board,reChoose=False,mct_rounds=2,depth=180)
    #print(move.point)
    #os.remove(hdf5)
    train_file=f5obj(hdf5)
    bot1.compile()
    while True:
        board = Board(size=4)
        game=Game(board)
        bot1_win=0
        bot2_win=0
        for i in range(4):
            if np.random.randint(2)==0:
                result,his=game.run_alpha_zero_train(play_b=bot1,play_w=bot2,isprint=True,mct_rounds=10,depth=100)
                if result==GameResult.bWin:
                    winner='b'
                    bot1_win+=1
                else:
                    winner='w'
                    bot2_win+=1
            else:
                result,his=game.run_alpha_zero_train(play_b=bot2,play_w=bot1,isprint=True,mct_rounds=10,depth=100)
                if result==GameResult.bWin:
                    winner='b'
                    bot2_win+=1
                else:
                    winner='w'
                    bot1_win+=1
            train_file.save_game(winner,his,boardSize=4)
            x,y=train_file.load_data()
            bot1.train(x,y)
            os.remove(hdf5)
            board.reset()
            game.reset(board)
        total=bot1_win+bot2_win
        bot1.unload_weights(weights)
        if binom_test(bot1_win, total, 0.5)<.05 and bot1_win/total>.5:  #bot1显著优秀
            bot2.load_weights(weights)
